#include "warray/warray_nd.h"
#include "test.h"
#include "custom_type.h"
#include "helper.h"

typedef boost::mpl::list< float,double,custom_type,std::complex<float>,std::complex<double> > test_types;

using warp::arraynd;
using warp::const_arraynd;

// We need this to suppress warnings when converting an int to a std::complex<float>
template<typename T> T cast_to(int i, const T&)
{
	return T(i);
}
template<typename T> std::complex<T> cast_to(int i, const std::complex<T>&)
{
	return std::complex<T>(T(i));
}

template<typename T, size_t N>
void test_constructors()
{
	{
		arraynd<T,N> x0;
		BOOST_CHECK_EQUAL( x0.dims(), extents_gen<N>(0) );
		BOOST_CHECK_EQUAL( x0.size(), size_t(0) );
		BOOST_CHECK_EQUAL( x0.begin(), x0.end() );
	}
	{
		const_arraynd<T,N> x0;
		BOOST_CHECK_EQUAL( x0.dims(), extents_gen<N>(0) );
		BOOST_CHECK_EQUAL( x0.size(), size_t(0) );
		BOOST_CHECK_EQUAL( x0.begin(), x0.end() );
	}
	{
		arraynd<T,N> x0(extents_gen<N>(3),T(1));
		BOOST_CHECK_EQUAL( x0.dims(), extents_gen<N>(3) );
		BOOST_CHECK_EQUAL( x0.size(), size_t(pow(3.0,double(N))) );
		BOOST_CHECK_NE( x0.c_array(), (T*)0 );
		for(typename arraynd<T,N>::iterator it=x0.begin();it!=x0.end();++it)
			BOOST_CHECK_EQUAL( *it, T(1) );
	}
	{
		const_arraynd<T,N> x0(extents_gen<N>(3),T(1));
		BOOST_CHECK_EQUAL( x0.dims(), extents_gen<N>(3) );
		BOOST_CHECK_EQUAL( x0.size(), size_t(pow(3.0,double(N))) );
		BOOST_CHECK_NE( x0.c_array(), (const T*)0 );
		for(typename const_arraynd<T,N>::const_iterator it=x0.begin();it!=x0.end();++it)
			BOOST_CHECK_EQUAL( *it, T(1) );
	}
	{
		arraynd<T,N> x0(extents_gen<N>(3));
		BOOST_CHECK_EQUAL( x0.dims(), extents_gen<N>(3) );
		BOOST_CHECK_EQUAL( x0.size(), size_t(pow(3.0,double(N))) );
		BOOST_CHECK_NE( x0.c_array(), (T*)0 );
	}
	{
		const T carray[5] = { T(1), T(2), T(3), T(4), T(5) };
		const_arraynd<T,N> x0(extents_gen<N-1>(1)[5],carray);
		BOOST_CHECK_EQUAL( x0.dims(), extents_gen<N-1>(1)[5] );
		BOOST_CHECK_EQUAL( x0.size(), size_t(5) );
		BOOST_CHECK_EQUAL( x0.c_array(), carray );
	}
	{
		T carray[5] = { T(1), T(2), T(3), T(4), T(5) };
		arraynd<T,N> x0(extents_gen<N-1>(1)[5],carray);
		BOOST_CHECK_EQUAL( x0.dims(), extents_gen<N-1>(1)[5] );
		BOOST_CHECK_EQUAL( x0.size(), size_t(5) );
		BOOST_CHECK_EQUAL( x0.c_array(), carray );
	}
	{
		const_arraynd<T,N> x0(extents_gen<N>(3),T(1));
		const_arraynd<T,N> x1(x0);
		BOOST_CHECK_EQUAL( x0.dims(), x1.dims() );
		BOOST_CHECK_EQUAL( x0.strides(), x1.strides() );
		BOOST_CHECK_EQUAL( x0.size(), x1.size() );
		BOOST_CHECK_EQUAL( x0.c_array(), x1.c_array() );
	}
	{
		arraynd<T,N> x0(extents_gen<N>(3),T(1));
		const_arraynd<T,N> x1(x0);
		BOOST_CHECK_EQUAL( x0.dims(), x1.dims() );
		BOOST_CHECK_EQUAL( x0.strides(), x1.strides() );
		BOOST_CHECK_EQUAL( x0.size(), x1.size() );
		BOOST_CHECK_EQUAL( x0.c_array(), x1.c_array() );
	}
	{
		arraynd<T,N> x0(extents_gen<N>(3),T(1));
		arraynd<T,N> x1(x0);
		BOOST_CHECK_EQUAL( x0.dims(), x1.dims() );
		BOOST_CHECK_EQUAL( x0.strides(), x1.strides() );
		BOOST_CHECK_EQUAL( x0.size(), x1.size() );
		BOOST_CHECK_EQUAL( x0.c_array(), x1.c_array() );
	}
}

BOOST_AUTO_TEST_CASE_TEMPLATE( constructors, T, test_types )
{
	test_constructors<T,1>();
	test_constructors<T,2>();
	test_constructors<T,3>();
}

template<typename T, size_t N>
void test_assignment()
{
	arraynd<T,N> x0(extents_gen<N>(3));
	const_arraynd<T,N> x1(extents_gen<N>(4),T(0));
	
	const_arraynd<T,N> x2;
	
	x2 = x0;
	BOOST_CHECK_EQUAL( x0.dims(), x2.dims() );
	BOOST_CHECK_EQUAL( x0.strides(), x2.strides() );
	BOOST_CHECK_EQUAL( x0.size(), x2.size() );
	BOOST_CHECK_EQUAL( x0.c_array(), x2.c_array() );
	
	x2 = x1;
	BOOST_CHECK_EQUAL( x1.dims(), x2.dims() );
	BOOST_CHECK_EQUAL( x1.strides(), x2.strides() );
	BOOST_CHECK_EQUAL( x1.size(), x2.size() );
	BOOST_CHECK_EQUAL( x1.c_array(), x2.c_array() );
	
	arraynd<T,N> x3;
	
	x3 = x0;
	BOOST_CHECK_EQUAL( x0.dims(), x3.dims() );
	BOOST_CHECK_EQUAL( x0.strides(), x3.strides() );
	BOOST_CHECK_EQUAL( x0.size(), x3.size() );
	BOOST_CHECK_EQUAL( x0.c_array(), x3.c_array() );
}

BOOST_AUTO_TEST_CASE_TEMPLATE( assignment, T, test_types )
{
	test_assignment<T,1>();
	test_assignment<T,2>();
	test_assignment<T,3>();
}

template<typename T, size_t N>
void test_iterators()
{
	arraynd<T,N> x0( extents_gen<N>(3) );
	T *data = x0.c_array();
	for(unsigned int i=0;i<x0.size();++i)
		data[i] = cast_to(i,T());
	
	BOOST_CHECK( x0.begin() < x0.end() );
	BOOST_CHECK_EQUAL( size_t(x0.end()-x0.begin()), x0.size() ); 
	
	unsigned int i=0;
	for(typename arraynd<T,N>::iterator it=x0.begin();it!=x0.end();++it,++i)
		BOOST_CHECK_EQUAL( *it, cast_to(i,T()) );
	
	const_arraynd<T,N> x1 = x0;
	
	BOOST_CHECK( x1.begin() < x1.end() );
	BOOST_CHECK_EQUAL( size_t(x1.end()-x1.begin()), x1.size() ); 
	
	i=0;
	for(typename const_arraynd<T,N>::const_iterator it=x1.begin();it!=x1.end();++it,++i)
		BOOST_CHECK_EQUAL( *it, cast_to(i,T()) );
}

BOOST_AUTO_TEST_CASE_TEMPLATE( iterators, T, test_types )
{
	test_iterators<T,1>();
	test_iterators<T,2>();
	test_iterators<T,3>();
}

BOOST_AUTO_TEST_CASE_TEMPLATE( array1d_access, T, test_types )
{
	arraynd<T,1> x0(warp::extents[3]);
	
	x0[0] = T(1);
	x0[1] = T(2);
	x0[2] = T(3);
	
	BOOST_CHECK_EQUAL( x0[0], T(1) );
	BOOST_CHECK_EQUAL( x0[1], T(2) );
	BOOST_CHECK_EQUAL( x0[2], T(3) );
	BOOST_CHECK_EQUAL( x0.c_array()[0], T(1) );
	BOOST_CHECK_EQUAL( x0.c_array()[1], T(2) );
	BOOST_CHECK_EQUAL( x0.c_array()[2], T(3) );
}

BOOST_AUTO_TEST_CASE_TEMPLATE( array2d_access, T, test_types )
{
	arraynd<T,2> x0(warp::extents[3][3]);
	
	int val=1;
	for(int i=0;i<3;++i)
		for(int j=0;j<3;++j)
			x0[i][j] = cast_to(val++,T());
		
	val = 1;
	for(int i=0;i<3;++i)
		for(int j=0;j<3;++j)
			BOOST_CHECK_EQUAL( x0[i][j], cast_to(val++,T()) );
	
	val = 1;
	for(int i=0;i<9;++i)
		BOOST_CHECK_EQUAL( x0.c_array()[i], cast_to(val++,T()) );
}

BOOST_AUTO_TEST_CASE_TEMPLATE( array3d_access, T, test_types )
{
	arraynd<T,3> x0(warp::extents[3][3][3]);
	
	int val=1;
	for(int i=0;i<3;++i)
		for(int j=0;j<3;++j)
			for(int k=0;k<3;++k)
				x0[i][j][k] = cast_to(val++,T());
		
	val = 1;
	for(int i=0;i<3;++i)
		for(int j=0;j<3;++j)
			for(int k=0;k<3;++k)
				BOOST_CHECK_EQUAL( x0[i][j][k], cast_to(val++,T()) );
	
	val = 1;
	for(int i=0;i<27;++i)
		BOOST_CHECK_EQUAL( x0.c_array()[i], cast_to(val++,T()) );
}

template<typename T, size_t N>
void test_copy()
{
	arraynd<T,N> x0(extents_gen<N>(3));
	fill_nonzero(x0);
	
	arraynd<T,N> x1 = x0.copy();
	BOOST_CHECK_EQUAL( x0.dims(), x1.dims() );
	BOOST_CHECK_EQUAL( x0.strides(), x1.strides() );
	BOOST_CHECK_EQUAL( x0.size(), x1.size() );
	BOOST_CHECK_NE( x0.c_array(), x1.c_array() );
	BOOST_CHECK_EQUAL_COLLECTIONS( x0.begin(), x0.end(), x1.begin(), x1.end() );
	
	const_arraynd<T,N> x2 = x0;
	x1 = x0.copy();
	BOOST_CHECK_EQUAL( x0.dims(), x1.dims() );
	BOOST_CHECK_EQUAL( x0.strides(), x1.strides() );
	BOOST_CHECK_EQUAL( x0.size(), x1.size() );
	BOOST_CHECK_NE( x0.c_array(), x1.c_array() );
	BOOST_CHECK_EQUAL_COLLECTIONS( x0.begin(), x0.end(), x1.begin(), x1.end() );
}

BOOST_AUTO_TEST_CASE_TEMPLATE( copy, T, test_types )
{
	test_copy<T,1>();
	test_copy<T,2>();
	test_copy<T,3>();
}

template<typename T, size_t N>
void test_inject()
{
	arraynd<T,N> x0(extents_gen<N>(3));
	fill_nonzero(x0);
	
	arraynd<T,N> x1(x0.dims());
	x1.inject(x0);
	BOOST_CHECK_EQUAL( x0.dims(), x1.dims() );
	BOOST_CHECK_EQUAL( x0.strides(), x1.strides() );
	BOOST_CHECK_EQUAL( x0.size(), x1.size() );
	BOOST_CHECK_NE( x0.c_array(), x1.c_array() );
	BOOST_CHECK_EQUAL_COLLECTIONS( x0.begin(), x0.end(), x1.begin(), x1.end() );
}

BOOST_AUTO_TEST_CASE_TEMPLATE( inject, T, test_types )
{
	test_inject<T,1>();
	test_inject<T,2>();
	test_inject<T,3>();
}


BOOST_AUTO_TEST_CASE_TEMPLATE( array2d_slice, T, test_types )
{
	{
		arraynd<T,2> x0( warp::extents[2][5] );
		arraynd<T,1> x1 = x0.slice( warp::indices[0] );
		arraynd<T,1> x2 = x0.slice( warp::indices[1] );
		
		BOOST_CHECK_EQUAL( x1.dim(warp::dim1), size_t(5) );
		BOOST_CHECK_EQUAL( x2.dim(warp::dim1), size_t(5) );
		BOOST_CHECK_EQUAL( x1.c_array(), x0.c_array() );
		BOOST_CHECK_EQUAL( x2.c_array(), x0.c_array()+5 );
	}
	
	{
		const_arraynd<T,2> x0( warp::extents[2][5], T(0) );
		const_arraynd<T,1> x1 = x0.slice( warp::indices[0] );
		const_arraynd<T,1> x2 = x0.slice( warp::indices[1] );
		
		BOOST_CHECK_EQUAL( x1.dim(warp::dim1), size_t(5) );
		BOOST_CHECK_EQUAL( x2.dim(warp::dim1), size_t(5) );
		BOOST_CHECK_EQUAL( x1.c_array(), x0.c_array() );
		BOOST_CHECK_EQUAL( x2.c_array(), x0.c_array()+5 );
	}
}

BOOST_AUTO_TEST_CASE_TEMPLATE( array3d_slice, T, test_types )
{
	{
		arraynd<T,3> x0( warp::extents[2][2][5] );
		arraynd<T,2> x1 = x0.slice( warp::indices[0] );
		arraynd<T,2> x2 = x0.slice( warp::indices[1] );
		arraynd<T,1> x3 = x0.slice( warp::indices[0][0] );
		arraynd<T,1> x4 = x0.slice( warp::indices[1][1] );
		
		BOOST_CHECK_EQUAL( x1.dim(warp::dim1), size_t(2) );
		BOOST_CHECK_EQUAL( x1.dim(warp::dim2), size_t(5) );
		BOOST_CHECK_EQUAL( x1.c_array(), x0.c_array() );
		BOOST_CHECK_EQUAL( x2.dim(warp::dim1), size_t(2) );
		BOOST_CHECK_EQUAL( x2.dim(warp::dim2), size_t(5) );
		BOOST_CHECK_EQUAL( x2.c_array(), x0.c_array()+10 );
		BOOST_CHECK_EQUAL( x3.dim(warp::dim1), size_t(5) );
		BOOST_CHECK_EQUAL( x3.c_array(), x0.c_array() );
		BOOST_CHECK_EQUAL( x4.dim(warp::dim1), size_t(5) );
		BOOST_CHECK_EQUAL( x4.c_array(), x0.c_array()+15 );
	}
	
	{
		const_arraynd<T,3> x0( warp::extents[2][2][5], T(0) );
		const_arraynd<T,2> x1 = x0.slice( warp::indices[0] );
		const_arraynd<T,2> x2 = x0.slice( warp::indices[1] );
		const_arraynd<T,1> x3 = x0.slice( warp::indices[0][0] );
		const_arraynd<T,1> x4 = x0.slice( warp::indices[1][1] );
		
		BOOST_CHECK_EQUAL( x1.dim(warp::dim1), size_t(2) );
		BOOST_CHECK_EQUAL( x1.dim(warp::dim2), size_t(5) );
		BOOST_CHECK_EQUAL( x1.c_array(), x0.c_array() );
		BOOST_CHECK_EQUAL( x2.dim(warp::dim1), size_t(2) );
		BOOST_CHECK_EQUAL( x2.dim(warp::dim2), size_t(5) );
		BOOST_CHECK_EQUAL( x2.c_array(), x0.c_array()+10 );
		BOOST_CHECK_EQUAL( x3.dim(warp::dim1), size_t(5) );
		BOOST_CHECK_EQUAL( x3.c_array(), x0.c_array() );
		BOOST_CHECK_EQUAL( x4.dim(warp::dim1), size_t(5) );
		BOOST_CHECK_EQUAL( x4.c_array(), x0.c_array()+15 );
	}
}

template<typename T, size_t N>
void test_strip()
{
	{
		arraynd<T,N> x0( extents_gen<N>(3) );
		arraynd<T,N> x1 = x0.strip(0,x0.dim(warp::dim1));
		arraynd<T,N> x2 = x0.strip(1,x0.dim(warp::dim1)-1);
		
		BOOST_CHECK_EQUAL( x0.dims(), x1.dims() );
		BOOST_CHECK_EQUAL( x0.strides(), x1.strides() );
		BOOST_CHECK_EQUAL( x0.size(), x1.size() );
		BOOST_CHECK_EQUAL( x0.c_array(), x1.c_array() );
	
		BOOST_CHECK_EQUAL( x2.dim(warp::dim1), size_t(1) );
		BOOST_CHECK_EQUAL( x0.strides(), x2.strides() );
		BOOST_CHECK_EQUAL( x0.size(), x2.size()*3 );
		BOOST_CHECK_EQUAL( x0.c_array()+extents_gen<N-1>(3).product(), x2.c_array() );
	}
	
	{
		const_arraynd<T,N> x0( extents_gen<N>(3), T(0) );
		const_arraynd<T,N> x1 = x0.strip(0,x0.dim(warp::dim1));
		const_arraynd<T,N> x2 = x0.strip(1,x0.dim(warp::dim1)-1);
		
		BOOST_CHECK_EQUAL( x0.dims(), x1.dims() );
		BOOST_CHECK_EQUAL( x0.strides(), x1.strides() );
		BOOST_CHECK_EQUAL( x0.size(), x1.size() );
		BOOST_CHECK_EQUAL( x0.c_array(), x1.c_array() );
	
		BOOST_CHECK_EQUAL( x2.dim(warp::dim1), size_t(1) );
		BOOST_CHECK_EQUAL( x0.strides(), x2.strides() );
		BOOST_CHECK_EQUAL( x0.size(), x2.size()*3 );
		BOOST_CHECK_EQUAL( x0.c_array()+extents_gen<N-1>(3).product(), x2.c_array() );
	}
}

BOOST_AUTO_TEST_CASE_TEMPLATE( strip, T, test_types )
{
	test_strip<T,1>();
	test_strip<T,2>();
	test_strip<T,3>();
}

